package com.zyx.flinkdemo.stream.operator.proc;

import com.zyx.flinkdemo.pojo.WordTs;
import com.zyx.flinkdemo.stream.utils.TimeUtils;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.TimeDomain;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.util.Collector;

import java.io.IOException;

/**
 * @author Yaxi.Zhang
 * @since 2021/6/29 09:20
 * desc: 基于事件时间
 */
public class EventTimerWithProcTimerDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<String> source = env.socketTextStream("localhost", 7777);
        //设置水印设置周期
        env.getConfig().setAutoWatermarkInterval(200);
        env.setParallelism(1);
        SingleOutputStreamOperator<WordTs> map = source.map(new MapFunction<String, WordTs>() {
            @Override
            public WordTs map(String value)  {
                String[] split = value.split(",");
                return new WordTs(split[0], Long.parseLong(split[1]));
            }
        });

        SingleOutputStreamOperator<WordTs> wordTsDs = map.assignTimestampsAndWatermarks(WatermarkStrategy
                .<WordTs>forMonotonousTimestamps()
                .withTimestampAssigner(new SerializableTimestampAssigner<WordTs>() {
                    @Override
                    public long extractTimestamp(WordTs element, long recordTimestamp) {
                        return element.getTs();
                    }
                }));

        wordTsDs
                .keyBy(WordTs::getWord)
                .process(new EventTimerWithProcessFunc())
                .print("wordCount >>> ");

        env.execute();
    }

    public static class EventTimerWithProcessFunc extends KeyedProcessFunction<String, WordTs, String> {
        // 记录统计个数
        private static ValueState<Long> countState;
        // 记录事件时间的触发时间
        private static ValueState<Long> eventTimer;

        @Override
        public void open(Configuration parameters) {
            // 初始化ValueState
            countState = getRuntimeContext().getState(new ValueStateDescriptor<>("count_state", Types.LONG));
            eventTimer = getRuntimeContext().getState(new ValueStateDescriptor<>("event_timer", Types.LONG));
        }

        @Override
        public void processElement(WordTs value, Context ctx, Collector<String> out) throws IOException {
            // 统计个数加一
            Long count = countState.value();
            if (count == null) {
                countState.update(1L);
            } else {
                countState.update(count + 1);
            }
            // 注册事件时间
            Long eventTimerTs = eventTimer.value();
            long currentWatermark = ctx.timerService().currentWatermark();
            if (eventTimerTs == null) {
                eventTimer.update(currentWatermark);
                ctx.timerService().registerEventTimeTimer(currentWatermark + 5000L);
            }
            // 注册处理时间
            long currentProcessingTime = ctx.timerService().currentProcessingTime();
            ctx.timerService().registerProcessingTimeTimer(currentProcessingTime + 5000L);
            String info = "watermark:" + TimeUtils.getFormatTimeStr(currentWatermark) + "proctime:" + TimeUtils.getFormatTimeStr(currentProcessingTime);
            out.collect("process >>> " + info);
        }

        @Override
        public void onTimer(long ts, OnTimerContext ctx, Collector<String> out) throws IOException {
            String info = "onTimer >>> key:" + ctx.getCurrentKey() + ", ts:" + TimeUtils.getFormatTimeStr(ts) + ", count:" + countState.value();
            out.collect(info);
            // 如果当前触发器 触发事件时间
            if (ctx.timeDomain() == TimeDomain.EVENT_TIME) {
                eventTimer.clear();
            }
        }
    }

}
